#ifndef ENDPOINT_STORE_H
#define ENDPOINT_STORE_H

#include "hccl_api.h"
#include <memory>
#include <vector>
#include <map>
#include <mutex>

namespace hixl {
    using EndpointHandle = void *;
    using ChannelHandle = void *;

    struct EndpointInfo {
        hccl::EndPoint endpoint;
        EndpointHandle handle;
    };

    class EndpointStore {
    public:
        EndpointStore();
        ~EndpointStore();

        hccl::HcclResult CreateEndpoint(const hccl::EndPoint &endpoint, EndpointHandle *endpoint_handle);
        hccl::HcclResult DestroyEndpoint(EndpointHandle endpoint_handle);

        // 新增：获取端点信息
        hccl::HcclResult GetEndpointInfo(EndpointHandle endpoint_handle, hccl::EndPoint *endpoint_info) const;

        // 新增：根据句柄查找端点信息
        const hccl::EndPoint* FindEndpoint(EndpointHandle endpoint_handle) const;

        hccl::HcclResult MemImport(EndpointHandle endpoint_handle, const void *memDesc, uint32_t descLen,
                                   hccl::HcommBuf *outBuf);

        hccl::HcclResult CreateChannel(EndpointHandle endpoint_handle, ChannelHandle *channel_handle);
        hccl::HcclResult DestroyChannel(EndpointHandle endpoint_handle, ChannelHandle channel_handle);

        std::vector<EndpointHandle> GetAllEndpointHandles() const;

    private:
        mutable std::mutex endpoints_mutex_;
        std::map<EndpointHandle, EndpointInfo> endpoints_info_;  // 存储端点信息和句柄的映射

        mutable std::mutex channels_mutex_;
        std::map<EndpointHandle, std::vector<ChannelHandle> > endpoint_channels_;
    };
} // namespace hixl

#endif // ENDPOINT_STORE_H